import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningDataModule


class CreateDataset(Dataset):
    def __init__(self, data_array):
        super().__init__()
        self.data_array = data_array

    def __getitem__(self, item):
        x = self.data_array[item]
        return torch.tensor(data=x, dtype=torch.float32)

    def __len__(self):
        return len(self.data_array)


class CreateDataset_LM(LightningDataModule):
    def __init__(self, data_array, batch_size, num_workers):
        super().__init__()
        self.train_dataset = None
        self.data_array = data_array
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        self.train_dataset = CreateDataset(data_array=self.data_array)

    def train_dataloader(self):
        return DataLoader(dataset=self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers,
                          shuffle=True)
